## code to run MSA-level MSMs and other models ##

msm <- function(dv1, dv2, dvp, num){
  
  print(paste0("MSA ", num))
  
  af.msa.reg <- remove_missing_from_mids(af.msa.mids, c(dv1, dv2))

  af.msa.t1t2 <- complete(af.msa.reg, "long", include = T)
  nrows <- nrow(af.msa.t1t2)/length(unique(af.msa.t1t2$.imp))
  
  print("COMPLETE OBS")
  print(nrow(af.msa.t1t2))
  
  vars.2022 <- paste(c("pop_total_2022","pop_density_2022",
                       "per_18t34_2022","per_35t64_2022",
                       "per_65over_2022","entropy_er_2022",
                       "per_child_2022", "per_nomove_2022",
                       "per_fb_2022","hownrt_2022",
                       "median_yr_str_2022","median_hhld_inc_2022",
                       "hhld_pov_rt_2022","unemr_2022",
                       "gini_2022"
  ), collapse = "+")
  
  vars.2009 <- paste(c("pop_total_2009","pop_density_2009",
                       "per_18t34_2009","per_35t64_2009",
                       "per_65over_2009","entropy_er_2009",
                       "per_child_2009", "per_nomove_2009",
                       "per_fb_2009","hownrt_2009",
                       "median_yr_str_2009","median_hhld_inc_2009",
                       "hhld_pov_rt_2009","unemr_2009",
                       "gini_2009"
  ), collapse = "+")
  
  fe.vars <- paste(c("pop_total","pop_density",
                     "per_18t34","per_35t64",
                     "per_65over","entropy_er",
                     "per_child", "per_nomove",
                     "per_fb","hownrt",
                     "median_yr_str","median_hhld_inc",
                     "hhld_pov_rt","unemr",
                     "gini","munis_tot",
                     "CBSA","tp"
  ), collapse = "+")
  
  ## regression formulas ##
  
  form1 <- paste0("treat_2022 ~ ",
                 dv1, " + treat_2006 + ",
                 vars.2022, " + munis_tot")
  
  form2 <- paste0("treat_2006 ~ ", 
                 vars.2009, " + munis_tot")
  
  form3a <- paste0(dv2," ~",
                  "thist + munis_tot")
  
  form3b <- paste0(dv2," ~",
                   "thist")
  
  form4 <- paste0(dv2," ~",
                  "zri_full_2022_st +",
                  "zri_full_2006_st +",
                  dv1," +",
                  vars.2022," +",
                  vars.2009, 
                  " + munis_tot")
  
  form5 <- paste0(dvp," ~",
                  "zri_full_st"," +",
                  fe.vars)
  
  msa.2022 <- with(af.msa.reg,
                   glm(as.formula(form1),
                       binomial(link = 'logit')))
  
  summary(pool(msa.2022))
  
  ## credit for how to do this: ##
  ## https://github.com/amices/mice/issues/82 ##
  
  ppout.2022 <- lapply(getfit(msa.2022), predict, type = "response", se.fit = TRUE)
  pps.2022 <- sapply(ppout.2022, `[[`, "fit")

  ## check predicted probabilities ##
  summary(pps.2022)
  
  ## convert the matrix of pps to a vector ##
  pps.2022.temp <- c()
  
  for(j in 1:ncol(pps.2022)) {
    pps.2022.temp <- c(pps.2022.temp, pps.2022[,j ])
  }
  
  prop.table(table(af.msa.t1t2$treat_2022))
  
  ## numerator for stabilized weights ##
  
  msa.2022.nm <- with(af.msa.reg,
                      glm(treat_2022 ~ 
                          treat_2006,
                      binomial(link = 'logit')))
  
  summary(pool(msa.2022.nm))
  
  ppout.2022.num <- lapply(getfit(msa.2022.nm), predict, type = "response", se.fit = TRUE)
  pps.2022.num <- sapply(ppout.2022.num, `[[`, "fit")
  
  ## check probabilities ##
  summary(pps.2022.num)
  
  pps.2022.num.rf <- c()
  for(j in 1:ncol(pps.2022.num)) {
    pps.2022.num.rf <- c(pps.2022.num.rf, pps.2022.num[,j ])
  }
  
  summary(pps.2022.num.rf)

  msa.2006 <- with(af.msa.reg,
                   glm(as.formula(form2),
                       binomial(link = 'logit')))
  
  summary(pool(msa.2006))
  
  ppout.2006 <- lapply(getfit(msa.2006), predict, type = "response", se.fit = TRUE)
  pps.2006 <- sapply(ppout.2006, `[[`, "fit")
  
  ## convert the matrix of pps to a vector ##
  pps.2006.temp <- c()
  
  for(j in 1:ncol(pps.2006)) {
    pps.2006.temp <- c(pps.2006.temp, pps.2006[,j ])
  }
  
  num06.calc <- af.msa.t1t2 %>%
    filter(.imp != 0) %>%
    group_by(.imp) %>%
    summarize(num = mean(treat_2006,na.rm=T)) %>%
    slice(rep(1:n(), each = nrows))
  
  num06 <- num06.calc$num
  
  print("LENGTHS - REG")
  print(length(pps.2022.temp))
  print(length(pps.2006.temp))

  ## concatenate NAs for .imp = 0 ##
  pps.2006.num.fin <- c(rep(NA,nrows),num06)
  pps.2022.fin <- c(rep(NA,nrows),pps.2022.temp)
  pps.2022.num.fin <- c(rep(NA,nrows),pps.2022.num.rf)
  pps.2006.fin <- c(rep(NA,nrows),pps.2006.temp)
  
  ## attach the pps and numerators and create weights ##
  
  af.msa.wts <- complete(af.msa.reg, 'long', include = TRUE) %>%
    mutate(thist = treat_2006 + treat_2022,
           num_t_2006 = pps.2006.num.fin,
           pps_2006 = pps.2006.fin,
           num_t_2022 = pps.2022.num.fin,
           pps_2022 = pps.2022.fin,
           num_2006 = ifelse(treat_2006 == 1, num_t_2006, 1-num_t_2006),
           den_2006 = ifelse(treat_2006 == 1, pps_2006, 1-pps_2006),
           num_2022 = ifelse(treat_2022 == 1, num_t_2022, 1-num_t_2022),
           den_2022 = ifelse(treat_2022 == 1, pps_2022, 1-pps_2022),
           swt_2006 = num_2006/den_2006,
           swt_2022 = num_2022/den_2022,
           swt_fin = swt_2006 * swt_2022)
  
  table(af.msa.wts$thist, useNA = "ifany")
  
  ## truncate the weights ##
  
  wts.cb.p5 <- quantile(af.msa.wts$swt_fin, 0.05, na.rm=T)
  wts.cb.p95 <- quantile(af.msa.wts$swt_fin,0.95, na.rm=T)
  
  af.msa.wts$swt_fin[af.msa.wts$swt_fin > wts.cb.p95] <- wts.cb.p95
  af.msa.wts$swt_fin[af.msa.wts$swt_fin < wts.cb.p5] <- wts.cb.p5
  summary(af.msa.wts$swt_fin)
  sd(af.msa.wts$swt_fin, na.rm=T)
  
  summary(af.msa.wts$swt_fin)
  summary(af.msa.wts$swt_fin[af.msa.wts$.imp == 0])
  summary(af.msa.wts$swt_fin[af.msa.wts$.imp != 0])
  
  ## assign the weights for export ##
  wts.fin <- af.msa.wts$swt_fin
  
  af.msa.mids.fin <- as.mids(af.msa.wts)
  
  ## msm ##
  
  msa.msm <- with(af.msa.mids.fin,
                  lm(as.formula(form3a),
                     weights=swt_fin))
  
  ## unadjusted ## 
  
  msa.noadj <- with(af.msa.mids.fin,
                    lm(as.formula(form3b)))
  
  ## adl ## 
  
  msa.adl <- with(af.msa.mids.fin,
                  lm(as.formula(form4)))
  
  ## fixed effects ## 

  msa.fe <- with(af.msa.fe.mids,
                 lm(as.formula(form5)))
  
  out.msm <- summary(pool(msa.msm))
  out.noadj <- summary(pool(msa.noadj))
  out.adl <- summary(pool(msa.adl))
  out.fe <- summary(pool(msa.fe))

  
  ## assess balance ##
  
  bal.t2.vars <- c("pop_total_2022","pop_density_2022",
                   "per_18t34_2022","per_35t64_2022",
                   "per_65over_2022","entropy_er_2022",
                   "per_child_2022", "per_nomove_2022",
                   "per_fb_2022","hownrt_2022",
                   "median_yr_str_2022","median_hhld_inc_2022",
                   "hhld_pov_rt_2022","unemr_2022",
                   "gini_2022")
  
  bal.t1.vars <- c("pop_total_2009","pop_density_2009",
                   "per_18t34_2009","per_35t64_2009",
                   "per_65over_2009","entropy_er_2009",
                   "per_child_2009", "per_nomove_2009",
                   "per_fb_2009","hownrt_2009",
                   "median_yr_str_2009","median_hhld_inc_2009",
                   "hhld_pov_rt_2009","unemr_2009",
                   "gini_2009")

  bal.t2 <- function(var){
    
    unwt2 <- with(af.msa.mids.fin,
                  lm(as.formula(paste0(var,
                                       " ~ treat_2022 + treat_2006 + munis_tot"))))
    
    out2a <- summary(pool(unwt2))
    
    wt2 <- with(af.msa.mids.fin,
                lm(as.formula(paste0(var,
                                     " ~ treat_2022 + treat_2006 + munis_tot")),
                   weights=swt_fin))
    
    out2b <- summary(pool(wt2))
    
    outres <- list(out2a,out2b)
    
    return(outres)
    
  }

  t2.bal <- lapply(bal.t2.vars, bal.t2)  
  
  bal.t1 <- function(var){
    
    unwt1 <- with(af.msa.mids.fin,
                  lm(as.formula(paste0(var,
                                       " ~ treat_2006 + munis_tot"))))
    
    out1a <- summary(pool(unwt1))
    
    
    wt1 <- with(af.msa.mids.fin,
                lm(as.formula(paste0(var,
                                     " ~ treat_2006 + munis_tot")),
                   weights=swt_fin))
    
    out1b <- summary(pool(wt1))
    
    outres <- list(out1a,out1b)
    
    return(outres)
    
  }
  
  t1.bal <- lapply(bal.t1.vars, bal.t1)  
  
  ## capture all output ## 
  
  res <- list(out.msm,
              wts.fin,
              out.noadj,
              out.adl,
              out.fe,
              t2.bal,
              t1.bal)
  
  return(res)
  
}

